import joblib
import os
import pandas as pd
import pyautogui
import numpy as np
from tools.interact_recoder_rel_move import test_rel_move, run_warm_up
import mouse
from bdtime import tt


input_dir_name = os.path.join('./tempdir', 'df')
output_dir_name = os.path.join('./tempdir', 'models')
output_df_dir_name = os.path.join('./tempdir', 'test_df')
if os.path.basename(os.getcwd()) in ['tests', 'tools']:
    input_dir_name = os.path.join('../tempdir', 'df')
    output_dir_name = os.path.join('../tempdir', 'models')
    output_df_dir_name = os.path.join('../tempdir', 'test_df')
    # input_dir_name = os.path.join('../..', input_dir_name)
    # output_dir_name = os.path.join('../..', output_dir_name)
    # output_df_dir_name = os.path.join('../..', output_df_dir_name)

os.makedirs(output_df_dir_name, exist_ok=True)

# model_path = os.path.join(output_dir_name, 'predict_rel_move.pkl')
# model_path = os.path.join(output_dir_name, 'predict_rel_move__10_50.pkl')
model_path = os.path.join(output_dir_name, 'predict_rel_move__400_600.pkl')

loaded_model = joblib.load(model_path)

threshold_range_ls = [
    [0, 10],
    [10, 25],
    [25, 50],
    [50, 100],
    [100, 200],
    [200, 400],
    [400, 600],
]
threshold_range_arr = np.array(threshold_range_ls)

model_ls = [joblib.load(os.path.join(output_dir_name, f'predict_rel_move__{ls_i[0]}_{ls_i[1]}.pkl')) for ls_i in threshold_range_ls]

screen_width, screen_height = pyautogui.size()

def predict_delta_by_dx_ls(ls, warm_up=True):
    if isinstance(ls, (int, float)):
        ls = [ls]
    _ls = []
    for i in range(len(ls)):
        o = ls[i]
        if i == 0 and not warm_up:
            idx = 0
        else:
            idx = 1
        _ls.append([idx, o])

    arr = np.array(ls) / 1.8
    max_value = np.max(np.abs(arr))

    distances = np.sum(np.abs(threshold_range_arr - max_value), axis=1)
    min_index = np.where(distances == min(distances))[0][0]
    threshold_range_i = threshold_range_ls[min_index]
    print(f'=== predict_delta_by_dx_ls --- max_value: {max_value}, threshold_range_i:', threshold_range_i)
    most_suitable_model = model_ls[min_index]

    res = np.round(most_suitable_model.predict(_ls)).astype(int)
    return res


def run_dx_ls(dx_ls, _interval=0, _init_position=None, show=True):
    if _init_position is None:
        _init_position = mouse.get_position()[0]
    else:
        mouse.move(_init_position, screen_height // 2)

    columns = ['i', 'dx', 'delta_0', 'rate_td_dx']
    df = pd.DataFrame([], columns=columns)

    total_i = len(dx_ls)
    for i in range(len(dx_ls)):
        dx = dx_ls[i]
        if tt.stop_alt('x'):
            raise EOFError("stop!!")

        print(f'\n=== i: {i} / {total_i}')

        if _interval:
            tt.sleep(_interval)

        delta_0, delta_1, delta_mean, rate_d0_dm, rate_dx_dm, total_dx, rate_td_dx, cost_t = test_rel_move(dx)

        # row = pd.DataFrame([[i, dx, delta_0, delta_1, delta_mean, rate_d0_dm, rate_dx_dm, total_dx]], columns=columns)
        row = pd.DataFrame([[i, dx, delta_0, rate_td_dx]], columns=columns)

        df = pd.concat([df, row], ignore_index=True)
    end_position = mouse.get_position()[0]
    total_dx = end_position - _init_position

    if show:
        print(f'\n~~~ init_position: {_init_position} --> end_position:', end_position, ', total_dx:', total_dx)
    return df


def statistic_df(df, loaded_model, warm_up=True, show=True):
    from sklearn.metrics import mean_squared_error

    df['_i'] = 1 if warm_up else df['i'].apply(lambda x: 0 if x == 0 else 1)
    df['predict'] = df[['_i', 'delta_0']].apply(lambda x: int(round(loaded_model.predict([[x['_i'], x['delta_0']]])[0])), axis=1)
    df['error_v'] = df[['dx', 'predict']].apply(lambda x: x['dx'] - x['predict'], axis=1)
    df['error_p'] = df[['dx', 'predict']].apply(lambda x: round(((x['dx'] - x['predict']) / x['dx']), 4), axis=1)

    if show:
        print(df)

    compute_metric = mean_squared_error
    metric_mse = compute_metric(df['dx'], df['predict'])

    from tools.lol.my_mse import mean_percentage_error
    compute_metric = mean_percentage_error
    metric_mpe = compute_metric(df['dx'], df['predict'])

    total_move = sum(df['dx'])
    accumulate_error_value = total_move - sum(df['predict'])
    accumulate_error_percent = accumulate_error_value / total_move if total_move else 0

    if show:
        print(
            f'total_move: {total_move},  accumulate_error_value: {accumulate_error_value: g}, accumulate_error_percent: {accumulate_error_percent: .3%}')

        print(f"~~~ mean_squared_error: {metric_mse: g}")
        print(f"~~~ mean_percentage_error: {metric_mpe:.2%}")
    metric_dc = {
        'metric_mse': metric_mse,
        'metric_mpe': metric_mpe,
        'total_move': total_move,
        'accumulate_error_value': accumulate_error_value,
        'accumulate_error_percent': accumulate_error_percent,
    }

    return df, metric_dc


if __name__ == '__main__':
    # _init_position = mouse.get_position()[0]

    # dx_ls = [100, -100, 50, -50, 40, -40, 30, -30, 20, -20, 10, -10]

    # dx_range = 20
    # dx_length = 10
    # dx_ls = np.random.randint(-dx_range, dx_range, dx_length)
    split_times = 2
    # duplicate_times = 16
    duplicate_times = 2
    dx_ls = [screen_width // split_times] * duplicate_times + [-screen_width // split_times] * duplicate_times

    _init_position = 1000
    _interval = 0.5
    # _interval = 0

    use_predict = True
    # use_predict = False

    warm_up = True

    if warm_up:
        run_warm_up()

    if use_predict:
        dx_ls = predict_delta_by_dx_ls(dx_ls)
        duplicate_times *= 234

    df = run_dx_ls(dx_ls, _interval=_interval,  _init_position=_init_position)

    tt.sleep(1)
    import keyboard
    keyboard.press_and_release('3')
    tt.sleep(2)
    keyboard.press_and_release('4')

    tt.sleep(1)

    df = run_dx_ls(dx_ls, _interval=_interval,  _init_position=_init_position)

    df, metric_dc = statistic_df(df, loaded_model)
